package com.ztorn.cdc.pulsar;

import com.ztorn.common.util.JSONUtil;
import com.fasterxml.jackson.databind.node.ObjectNode;
import lombok.extern.slf4j.Slf4j;
import org.apache.flink.api.common.serialization.SerializationSchema;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.functions.sink.RichSinkFunction;
import org.apache.pulsar.client.api.*;

import java.util.concurrent.CompletableFuture;
import java.util.concurrent.TimeUnit;

@Slf4j
public class PulsarSinkFunction<T> extends RichSinkFunction<T> {

    private final String topic;
    private final String serviceUrl;
    private SerializationSchema<T> serializationSchema;

    Producer<byte[]> producer = null;

    public PulsarSinkFunction(String topic, String serviceUrl, SerializationSchema<T> serializationSchema){
        this.topic = topic;
        this.serviceUrl = serviceUrl;
        this.serializationSchema = serializationSchema;
    }

    @Override
    public void open(Configuration parameters) throws Exception {
        PulsarClient client = PulsarClient.builder().serviceUrl(serviceUrl).build();
        producer = client.newProducer()
                .topic(topic)
                .compressionType(CompressionType.LZ4)
                .sendTimeout(0, TimeUnit.SECONDS)
                .enableBatching(true)
                .batchingMaxPublishDelay(10, TimeUnit.MILLISECONDS)
                .batchingMaxMessages(1000)
                .maxPendingMessages(1000)
                .blockIfQueueFull(true)
                .roundRobinRouterBatchingPartitionSwitchFrequency(10)
                .batcherBuilder(BatcherBuilder.DEFAULT)
                .create();
    }

    @Override
    public void invoke(T value, Context context) throws Exception {
        byte[] serializeValue = serializationSchema.serialize(value);
        String strValue = new String(serializeValue);
        CompletableFuture<MessageId> future = producer.newMessage()
                .key(getKey(strValue))
                .value(serializeValue).sendAsync();
        future.handle((v, ex) -> {
            if (ex == null) {
                log.info("发送消息成功, key:{}, msg: {}", getKey(strValue), strValue);
            } else {
                log.error("发送消息失败, key:{}, msg: {}", getKey(strValue), strValue);
            }
            return null;
        });
        future.join();
        log.info("发送消息完成, key:{}, msg: {}", getKey(strValue), strValue);
    }

    @Override
    public void close() throws Exception {

    }

    public String getKey(String strValue) {
        ObjectNode jsonNodes = JSONUtil.parseObject(strValue);
        String key = String.valueOf(jsonNodes.get("key"));
        return key == null ? "" : key;
    }
}
